# states/context_partition_manager.py（上下文分区管理器）
from typing import Dict, List, Optional, Any, Set
from dataclasses import dataclass
from enum import Enum
import asyncio
import json
from datetime import datetime, timedelta

from utils.logger import Logger
from states.langgraph_workflow import TaskType

# 本地定义AgentType枚举以避免导入错误
class AgentType(Enum):
    TEXT_QA = "text_qa"
    DOC_QA = "doc_qa"
    MULTI_MODAL = "multi_modal"
    OCR = "ocr"
    ORCHESTRATOR = "orchestrator"
from core.cache.multilevel_cache import MultiLevelCache, CacheConfig
from middleware.redis_adapter import get_redis_adapter


class ContextPartitionType(Enum):
    """
    上下文分区类型枚举
    用于标识不同类型的记忆空间分区
    """
    TEXT_QA = "text_qa"               # 文本问答上下文
    CODE_GENERATION = "code_generation" # 代码生成上下文
    DOC_QA = "doc_qa"                 # 文档问答上下文
    MULTIMODAL = "multimodal"          # 多模态上下文
    CHAT_HISTORY = "chat_history"       # 聊天历史上下文
    TOOL_CALL = "tool_call"             # 工具调用上下文
    ORCHESTRATOR = "orchestrator"       # 协调器上下文


class PartitionStrategy(Enum):
    """
    分区策略枚举
    定义如何管理不同分区的记忆空间
    """
    ISOLATED = "isolated"           # 完全隔离（不同分区完全独立）
    PARTIAL_SHARED = "partial_shared" # 部分共享（部分通用信息共享）
    HIERARCHICAL = "hierarchical"     # 分层共享（基于层次结构共享）


@dataclass
class ContextEntry:
    """
    上下文条目数据类
    表示单个上下文记忆项
    """
    entry_id: str           # 条目唯一标识
    content: Dict[str, Any] # 上下文内容
    created_at: float       # 创建时间戳
    updated_at: float       # 更新时间戳
    importance: float = 0.5 # 重要性评分（0-1）
    tags: List[str] = None  # 标签列表
    source_partition: Optional[str] = None  # 来源分区

    def __post_init__(self):
        if self.tags is None:
            self.tags = []


@dataclass
class PartitionConfig:
    """
    分区配置数据类
    定义每个分区的配置参数
    """
    partition_type: ContextPartitionType
    max_entries: int = 100           # 最大条目数
    ttl_seconds: int = 3600          # 过期时间（秒）
    importance_threshold: float = 0.3  # 重要性阈值
    enable_sharing: bool = False      # 是否允许共享
    shared_tags: Set[str] = None      # 允许共享的标签

    def __post_init__(self):
        if self.shared_tags is None:
            self.shared_tags = set()


class ContextPartitionManager:
    """
    上下文分区管理器
    为不同任务类型维护独立记忆空间，并提供分区管理功能
    """
    
    def __init__(self, cache_config: Optional[CacheConfig] = None):
        self.logger = Logger.get_logger(agent_name="context_partition_manager")
        self.cache_config = cache_config or CacheConfig()
        self.cache = MultiLevelCache(self.cache_config)
        self.redis_client = get_redis_adapter()
        self.lock = asyncio.Lock()
        
        # 分区配置字典
        self.partition_configs: Dict[ContextPartitionType, PartitionConfig] = {
            ContextPartitionType.TEXT_QA: PartitionConfig(
                partition_type=ContextPartitionType.TEXT_QA,
                max_entries=50,
                ttl_seconds=1800,
                enable_sharing=True,
                shared_tags={"common_knowledge", "general_facts"}
            ),
            ContextPartitionType.CODE_GENERATION: PartitionConfig(
                partition_type=ContextPartitionType.CODE_GENERATION,
                max_entries=80,
                ttl_seconds=2400,
                importance_threshold=0.4,
                enable_sharing=False
            ),
            ContextPartitionType.DOC_QA: PartitionConfig(
                partition_type=ContextPartitionType.DOC_QA,
                max_entries=120,
                ttl_seconds=3600,
                enable_sharing=True,
                shared_tags={"document_metadata", "section_overview"}
            ),
            ContextPartitionType.MULTIMODAL: PartitionConfig(
                partition_type=ContextPartitionType.MULTIMODAL,
                max_entries=60,
                ttl_seconds=1800,
                importance_threshold=0.45,
                enable_sharing=True,
                shared_tags={"visual_summaries", "content_overview"}
            ),
            ContextPartitionType.CHAT_HISTORY: PartitionConfig(
                partition_type=ContextPartitionType.CHAT_HISTORY,
                max_entries=200,
                ttl_seconds=7200,
                importance_threshold=0.2
            ),
            ContextPartitionType.TOOL_CALL: PartitionConfig(
                partition_type=ContextPartitionType.TOOL_CALL,
                max_entries=70,
                ttl_seconds=1200,
                enable_sharing=True,
                shared_tags={"tool_results", "execution_status"}
            ),
            ContextPartitionType.ORCHESTRATOR: PartitionConfig(
                partition_type=ContextPartitionType.ORCHESTRATOR,
                max_entries=150,
                ttl_seconds=4800,
                enable_sharing=True,
                shared_tags={"task_coordination", "agent_assignments"}
            )
        }
        
        # 启动清理任务 - 使用同步方式避免事件循环问题
        import threading
        def cleanup_thread():
            """同步清理线程"""
            try:
                with asyncio.Runner() as runner:
                    runner.run(self._cleanup_task())
            except Exception as e:
                self.logger.error(f"上下文分区清理任务失败: {str(e)}")
        
        # 启动清理线程
        threading.Thread(target=cleanup_thread, daemon=True).start()
        self.logger.info("上下文分区管理器初始化完成")
    
    def _generate_partition_key(self, partition_type: ContextPartitionType, user_id: Optional[str] = None) -> str:
        """
        生成分区键
        
        Args:
            partition_type: 分区类型
            user_id: 用户ID（可选，用于用户级别的分区隔离）
            
        Returns:
            分区键字符串
        """
        if user_id:
            return f"context:partition:{partition_type.value}:{user_id}"
        return f"context:partition:{partition_type.value}"
    
    def _generate_entry_key(self, partition_type: ContextPartitionType, entry_id: str, user_id: Optional[str] = None) -> str:
        """
        生成上下文条目键
        
        Args:
            partition_type: 分区类型
            entry_id: 条目ID
            user_id: 用户ID（可选）
            
        Returns:
            条目键字符串
        """
        return f"{self._generate_partition_key(partition_type, user_id)}:entry:{entry_id}"
    
    async def add_context_entry(
        self, 
        partition_type: ContextPartitionType,
        content: Dict[str, Any],
        importance: float = 0.5,
        tags: Optional[List[str]] = None,
        user_id: Optional[str] = None
    ) -> str:
        """
        添加上下文条目到指定分区
        
        Args:
            partition_type: 目标分区类型
            content: 上下文内容
            importance: 重要性评分（0-1）
            tags: 标签列表
            user_id: 用户ID（可选）
            
        Returns:
            创建的条目ID
        """
        async with self.lock:
            # 验证分区配置
            if partition_type not in self.partition_configs:
                raise ValueError(f"未知的分区类型: {partition_type}")
            
            config = self.partition_configs[partition_type]
            
            # 创建条目
            entry_id = f"{partition_type.value}_{int(datetime.now().timestamp())}_{hash(json.dumps(content))}"
            entry = ContextEntry(
                entry_id=entry_id,
                content=content,
                created_at=datetime.now().timestamp(),
                updated_at=datetime.now().timestamp(),
                importance=importance,
                tags=tags or [],
                source_partition=partition_type.value
            )
            
            # 检查是否达到最大条目数
            partition_key = self._generate_partition_key(partition_type, user_id)
            entry_ids = await self._get_partition_entries(partition_key)
            
            if len(entry_ids) >= config.max_entries:
                # 删除最不重要的条目
                await self._remove_least_important(partition_type, user_id, config.max_entries - 1)
            
            # 存储条目
            entry_key = self._generate_entry_key(partition_type, entry_id, user_id)
            await self.cache.set(entry_key, entry.__dict__, ttl=config.ttl_seconds)
            
            # 更新分区索引
            entry_ids.append(entry_id)
            await self.redis_client.sadd(partition_key, *entry_ids)
            await self.redis_client.expire(partition_key, config.ttl_seconds)
            
            # 处理共享条目
            if config.enable_sharing and any(tag in config.shared_tags for tag in entry.tags):
                await self._handle_shared_entry(entry, user_id)
            
            self.logger.debug(f"已添加上下文条目到分区 {partition_type.value}: {entry_id}")
            return entry_id
    
    async def get_context_entry(
        self, 
        partition_type: ContextPartitionType,
        entry_id: str,
        user_id: Optional[str] = None
    ) -> Optional[ContextEntry]:
        """
        获取指定分区中的上下文条目
        
        Args:
            partition_type: 分区类型
            entry_id: 条目ID
            user_id: 用户ID（可选）
            
        Returns:
            上下文条目或None
        """
        entry_key = self._generate_entry_key(partition_type, entry_id, user_id)
        entry_data = await self.cache.get(entry_key)
        
        if entry_data:
            return ContextEntry(**entry_data)
        return None
    
    async def get_partition_context(
        self, 
        partition_type: ContextPartitionType,
        limit: Optional[int] = None,
        importance_threshold: Optional[float] = None,
        tags: Optional[List[str]] = None,
        user_id: Optional[str] = None
    ) -> List[ContextEntry]:
        """
        获取分区的上下文条目列表
        
        Args:
            partition_type: 分区类型
            limit: 返回条目数量限制
            importance_threshold: 重要性阈值
            tags: 过滤标签
            user_id: 用户ID（可选）
            
        Returns:
            上下文条目列表
        """
        if partition_type not in self.partition_configs:
            raise ValueError(f"未知的分区类型: {partition_type}")
        
        config = self.partition_configs[partition_type]
        partition_key = self._generate_partition_key(partition_type, user_id)
        entry_ids = await self._get_partition_entries(partition_key)
        
        # 获取所有条目
        entries = []
        for entry_id in entry_ids:
            entry = await self.get_context_entry(partition_type, entry_id, user_id)
            if entry:
                entries.append(entry)
        
        # 应用过滤条件
        if importance_threshold is None:
            importance_threshold = config.importance_threshold
        
        filtered_entries = [
            entry for entry in entries 
            if entry.importance >= importance_threshold
        ]
        
        if tags:
            filtered_entries = [
                entry for entry in filtered_entries
                if any(tag in entry.tags for tag in tags)
            ]
        
        # 按重要性和时间排序
        filtered_entries.sort(
            key=lambda x: (x.importance, x.updated_at),
            reverse=True
        )
        
        # 返回限制数量
        if limit:
            return filtered_entries[:limit]
        
        return filtered_entries
    
    async def get_context_for_task(
        self, 
        task_type: TaskType,
        limit: int = 20,
        include_shared: bool = True,
        user_id: Optional[str] = None
    ) -> List[ContextEntry]:
        """
        获取特定任务类型的上下文，自动映射到相应分区
        
        Args:
            task_type: 任务类型
            limit: 返回条目数量限制
            include_shared: 是否包含共享上下文
            user_id: 用户ID（可选）
            
        Returns:
            上下文条目列表
        """
        # 任务类型到分区类型的映射
        task_partition_map = {
            TaskType.TEXT_QA: ContextPartitionType.TEXT_QA,
            TaskType.CODE_GENERATION: ContextPartitionType.CODE_GENERATION,
            TaskType.DOC_QA: ContextPartitionType.DOC_QA,
            TaskType.MULTIMODAL: ContextPartitionType.MULTIMODAL,
            TaskType.TOOL_CALL: ContextPartitionType.TOOL_CALL,
            TaskType.ORCHESTRATION: ContextPartitionType.ORCHESTRATOR
        }
        
        # 获取主分区
        partition_type = task_partition_map.get(
            task_type, 
            ContextPartitionType.CHAT_HISTORY  # 默认使用聊天历史分区
        )
        
        # 获取主分区上下文
        entries = await self.get_partition_context(partition_type, limit=limit, user_id=user_id)
        
        # 添加共享上下文
        if include_shared:
            shared_entries = await self._get_shared_context_for_task(
                task_type, 
                max_count=limit // 2,  # 共享上下文最多占一半
                user_id=user_id
            )
            
            # 合并并去重
            entry_ids = {entry.entry_id for entry in entries}
            for shared_entry in shared_entries:
                if shared_entry.entry_id not in entry_ids:
                    entries.append(shared_entry)
                    entry_ids.add(shared_entry.entry_id)
            
            # 重新排序
            entries.sort(
                key=lambda x: (x.importance, x.updated_at),
                reverse=True
            )
            entries = entries[:limit]
        
        return entries
    
    async def update_context_entry(
        self, 
        partition_type: ContextPartitionType,
        entry_id: str,
        content: Optional[Dict[str, Any]] = None,
        importance: Optional[float] = None,
        tags: Optional[List[str]] = None,
        user_id: Optional[str] = None
    ) -> bool:
        """
        更新上下文条目
        
        Args:
            partition_type: 分区类型
            entry_id: 条目ID
            content: 更新的内容
            importance: 更新的重要性评分
            tags: 更新的标签列表
            user_id: 用户ID（可选）
            
        Returns:
            是否更新成功
        """
        entry = await self.get_context_entry(partition_type, entry_id, user_id)
        if not entry:
            return False
        
        # 更新字段
        if content is not None:
            entry.content = content
        if importance is not None:
            entry.importance = importance
        if tags is not None:
            entry.tags = tags
        entry.updated_at = datetime.now().timestamp()
        
        # 保存更新
        entry_key = self._generate_entry_key(partition_type, entry_id, user_id)
        config = self.partition_configs[partition_type]
        result = await self.cache.set(entry_key, entry.__dict__, ttl=config.ttl_seconds)
        
        # 更新共享状态
        if config.enable_sharing and any(tag in config.shared_tags for tag in entry.tags):
            await self._handle_shared_entry(entry, user_id)
        
        return result
    
    async def delete_context_entry(
        self, 
        partition_type: ContextPartitionType,
        entry_id: str,
        user_id: Optional[str] = None
    ) -> bool:
        """
        删除上下文条目
        
        Args:
            partition_type: 分区类型
            entry_id: 条目ID
            user_id: 用户ID（可选）
            
        Returns:
            是否删除成功
        """
        async with self.lock:
            # 删除条目
            entry_key = self._generate_entry_key(partition_type, entry_id, user_id)
            delete_result = await self.cache.delete(entry_key)
            
            # 从分区索引中移除
            partition_key = self._generate_partition_key(partition_type, user_id)
            await self.redis_client.srem(partition_key, entry_id)
            
            # 删除共享条目
            await self._remove_shared_entry(entry_id, user_id)
            
            return delete_result
    
    async def clear_partition(
        self, 
        partition_type: ContextPartitionType,
        user_id: Optional[str] = None
    ) -> bool:
        """
        清空指定分区
        
        Args:
            partition_type: 分区类型
            user_id: 用户ID（可选）
            
        Returns:
            是否清空成功
        """
        async with self.lock:
            partition_key = self._generate_partition_key(partition_type, user_id)
            entry_ids = await self._get_partition_entries(partition_key)
            
            # 删除所有条目
            for entry_id in entry_ids:
                entry_key = self._generate_entry_key(partition_type, entry_id, user_id)
                await self.cache.delete(entry_key)
                await self._remove_shared_entry(entry_id, user_id)
            
            # 清空分区索引
            await self.redis_client.delete(partition_key)
            self.logger.info(f"已清空分区: {partition_type.value}")
            return True
    
    async def _get_partition_entries(self, partition_key: str) -> List[str]:
        """
        获取分区中的所有条目ID
        
        Args:
            partition_key: 分区键
            
        Returns:
            条目ID列表
        """
        try:
            return [item.decode('utf-8') if isinstance(item, bytes) else item 
                    for item in await self.redis_client.smembers(partition_key)]
        except Exception as e:
            self.logger.error(f"获取分区条目失败: {str(e)}")
            return []
    
    async def _remove_least_important(
        self, 
        partition_type: ContextPartitionType,
        user_id: Optional[str],
        keep_count: int
    ) -> None:
        """
        删除最不重要的条目，保留指定数量
        
        Args:
            partition_type: 分区类型
            user_id: 用户ID
            keep_count: 保留条目数量
        """
        entries = await self.get_partition_context(partition_type, user_id=user_id)
        
        # 保留最重要的条目
        if len(entries) > keep_count:
            entries_to_delete = entries[keep_count:]
            for entry in entries_to_delete:
                await self.delete_context_entry(partition_type, entry.entry_id, user_id)
    
    async def _handle_shared_entry(self, entry: ContextEntry, user_id: Optional[str]) -> None:
        """
        处理共享条目
        
        Args:
            entry: 上下文条目
            user_id: 用户ID
        """
        # 为共享条目创建特殊的键
        shared_key = f"context:shared:entry:{entry.entry_id}"
        
        # 存储共享条目（内容简化，仅包含共享信息）
        shared_content = {
            "entry_id": entry.entry_id,
            "content": entry.content,
            "importance": entry.importance,
            "tags": entry.tags,
            "source_partition": entry.source_partition,
            "shared_at": datetime.now().timestamp()
        }
        
        # 设置较长的TTL用于共享条目
        await self.redis_client.setex(
            shared_key,
            timedelta(hours=24),  # 共享条目保留24小时
            json.dumps(shared_content)
        )
        
        # 添加到共享索引
        await self.redis_client.sadd("context:shared:index", entry.entry_id)
    
    async def _remove_shared_entry(self, entry_id: str, user_id: Optional[str]) -> None:
        """
        移除共享条目
        
        Args:
            entry_id: 条目ID
            user_id: 用户ID
        """
        shared_key = f"context:shared:entry:{entry_id}"
        await self.redis_client.delete(shared_key)
        await self.redis_client.srem("context:shared:index", entry_id)
    
    async def _get_shared_context_for_task(
        self, 
        task_type: TaskType,
        max_count: int,
        user_id: Optional[str]
    ) -> List[ContextEntry]:
        """
        获取任务相关的共享上下文
        
        Args:
            task_type: 任务类型
            max_count: 最大数量
            user_id: 用户ID
            
        Returns:
            共享上下文条目列表
        """
        # 任务类型对应的相关标签
        task_related_tags = {
            TaskType.TEXT_QA: ["common_knowledge", "general_facts"],
            TaskType.CODE_GENERATION: ["tool_results", "execution_status"],
            TaskType.DOC_QA: ["document_metadata", "content_overview"],
            TaskType.MULTIMODAL: ["visual_summaries", "content_overview"],
            TaskType.TOOL_CALL: ["execution_status", "agent_assignments"],
            TaskType.ORCHESTRATION: ["task_coordination", "agent_assignments"]
        }
        
        related_tags = task_related_tags.get(task_type, [])
        if not related_tags:
            return []
        
        # 获取所有共享条目ID
        shared_entry_ids = await self.redis_client.smembers("context:shared:index")
        shared_entries = []
        
        for entry_id_bytes in shared_entry_ids:
            entry_id = entry_id_bytes.decode('utf-8') if isinstance(entry_id_bytes, bytes) else entry_id_bytes
            shared_key = f"context:shared:entry:{entry_id}"
            
            # 获取共享条目
            try:
                shared_data = await self.redis_client.get(shared_key)
                if shared_data:
                    data = json.loads(shared_data)
                    # 检查是否包含相关标签
                    if any(tag in data.get('tags', []) for tag in related_tags):
                        # 转换为ContextEntry
                        entry = ContextEntry(
                            entry_id=data['entry_id'],
                            content=data['content'],
                            created_at=data.get('shared_at', datetime.now().timestamp()),
                            updated_at=data.get('shared_at', datetime.now().timestamp()),
                            importance=data.get('importance', 0.5),
                            tags=data.get('tags', []),
                            source_partition=data.get('source_partition')
                        )
                        shared_entries.append(entry)
            except Exception as e:
                self.logger.error(f"获取共享条目失败 {entry_id}: {str(e)}")
        
        # 按重要性和时间排序
        shared_entries.sort(
            key=lambda x: (x.importance, x.updated_at),
            reverse=True
        )
        
        return shared_entries[:max_count]
    
    async def _cleanup_task(self):
        """
        定期清理过期条目和维护分区
        """
        while True:
            try:
                await asyncio.sleep(300)  # 每5分钟执行一次
                
                for partition_type, config in self.partition_configs.items():
                    # 清理每个分区
                    await self._cleanup_partition(partition_type, config)
                
                # 清理过期的共享条目
                await self._cleanup_shared_entries()
                
                self.logger.debug("上下文分区清理完成")
                
            except Exception as e:
                self.logger.error(f"上下文分区清理任务失败: {str(e)}")
    
    async def _cleanup_partition(self, partition_type: ContextPartitionType, config: PartitionConfig):
        """
        清理指定分区
        
        Args:
            partition_type: 分区类型
            config: 分区配置
        """
        try:
            # 获取所有分区键（包括用户级别的）
            pattern = f"context:partition:{partition_type.value}:*"
            partition_keys = await self.redis_client.keys(pattern)
            
            # 清理每个分区键
            for partition_key in partition_keys:
                # 获取并检查所有条目
                entry_ids = await self._get_partition_entries(partition_key)
                expired_count = 0
                
                for entry_id in entry_ids:
                    # 构建条目键
                    parts = partition_key.split(':')
                    user_id = parts[3] if len(parts) > 3 else None
                    entry_key = self._generate_entry_key(partition_type, entry_id, user_id)
                    
                    # 检查条目是否存在
                    if not await self.cache.get(entry_key):
                        # 从分区索引中移除不存在的条目
                        await self.redis_client.srem(partition_key, entry_id)
                        expired_count += 1
                
                if expired_count > 0:
                    self.logger.debug(f"清理分区 {partition_type.value} 中的 {expired_count} 个过期条目")
        
        except Exception as e:
            self.logger.error(f"清理分区 {partition_type.value} 失败: {str(e)}")
    
    async def _cleanup_shared_entries(self):
        """
        清理过期的共享条目
        """
        try:
            shared_entry_ids = await self.redis_client.smembers("context:shared:index")
            removed_count = 0
            
            for entry_id_bytes in shared_entry_ids:
                entry_id = entry_id_bytes.decode('utf-8') if isinstance(entry_id_bytes, bytes) else entry_id_bytes
                shared_key = f"context:shared:entry:{entry_id}"
                
                # 检查条目是否存在
                if not await self.redis_client.get(shared_key):
                    await self._remove_shared_entry(entry_id, None)
                    removed_count += 1
            
            if removed_count > 0:
                self.logger.debug(f"清理 {removed_count} 个过期的共享上下文条目")
                
        except Exception as e:
            self.logger.error(f"清理共享条目失败: {str(e)}")
    
    def get_partition_stats(self, partition_type: ContextPartitionType) -> Dict[str, Any]:
        """
        获取分区统计信息
        
        Args:
            partition_type: 分区类型
            
        Returns:
            统计信息字典
        """
        if partition_type not in self.partition_configs:
            raise ValueError(f"未知的分区类型: {partition_type}")
        
        config = self.partition_configs[partition_type]
        
        # 这里可以添加更多统计逻辑
        return {
            "partition_type": partition_type.value,
            "max_entries": config.max_entries,
            "ttl_seconds": config.ttl_seconds,
            "importance_threshold": config.importance_threshold,
            "enable_sharing": config.enable_sharing,
            "shared_tags": list(config.shared_tags)
        }


# 创建全局上下文分区管理器实例
context_partition_manager = ContextPartitionManager()
